import torch
import torch.nn as nn

from src.models.modules import LSTMEncoder, CNN_Extractor, Predictor, StationGatedFusion, BatchCombinedModel


class MetPolModel(nn.Module):
    def __init__(self, in_channels, local_in, local_out, global_in, global_out, input_size, hidden_size, meteor_hidden_size, fusion_hidden, **kwargs):
        super().__init__()
        
        self.cnn = CNN_Extractor(in_channels)
        self.gnn = BatchCombinedModel(local_in, local_out, global_in, global_out)
        self.lstm_local = LSTMEncoder(input_size, hidden_size)
        self.lstm_global = LSTMEncoder(input_size, hidden_size)
        self.fusion = StationGatedFusion(fusion_hidden)
        self.predictor = Predictor(meteor_hidden_size, hidden_size)
    
    def forward(self, met, met_time, tot_x, coord_info, map_order, month, day, hour, time_shift):
        met_info = self.cnn(met, met_time)
        global_output, local_total = self.gnn(tot_x, coord_info)
        global_fused = global_output[:, :, map_order, :]
        local_total = self.lstm_local(local_total, month, day, hour)
        global_fused = self.lstm_global(global_fused, month, day, hour)
        pol_info = self.fusion(local_total, global_fused)
        out, _ = self.predictor(met_info, pol_info, time_shift)
        return out
